import os
import json
from omics_bert_label_encoding.label.make_labels_charls import MakeLabelsCharls

class CheckLabelFilesCharls:
    def __init__(self, check_path: str, expected_line_count: int, regenerate: bool, dataset_path: str, dta_file_path: str, api_key: str = None, instance_mapping: dict = None):
        self.check_path = check_path
        self.expected_line_count = expected_line_count
        self.regenerate = regenerate
        self.dataset_path = dataset_path
        self.dta_file_path = dta_file_path
        self.api_key = api_key
        self.instance_mapping = instance_mapping

    def check_and_regenerate(self):
        if not os.path.exists(self.check_path):
            print(f"Directory not found: {self.check_path}")
            return

        files = [f for f in os.listdir(self.check_path) if f.endswith('.json')]
        if not files:
            print("No JSON files found to check.")
            return

        for file in files:
            file_path = os.path.join(self.check_path, file)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    if not isinstance(data, list) or len(data) < self.expected_line_count:
                        print(f"File {file} has fewer than {self.expected_line_count} lines. Regenerating...")
                        if self.regenerate:
                            self.regenerate_label(file, file_path)
            except json.JSONDecodeError:
                print(f"File {file} is not a valid JSON. Regenerating...")
                if self.regenerate:
                    self.regenerate_label(file, file_path)
            except Exception as e:
                print(f"An error occurred with file {file}: {e}")

    def regenerate_label(self, file_name: str, file_path: str):
        column_name = os.path.splitext(file_name)[0]
        print(f"Regenerating label for {column_name}...")
        try:
            # Temporarily remove the problematic file to avoid conflicts
            if os.path.exists(file_path):
                os.remove(file_path)
            
            # Initialize the label maker and call the specific processing function
            label_maker = MakeLabelsCharls(self.dataset_path, os.path.dirname(self.check_path), self.dta_file_path, self.api_key, self.instance_mapping)
            label_maker.process_column_name(column_name)
            print(f"Successfully regenerated label for {column_name}.")
        except Exception as e:
            print(f"Failed to regenerate label for {column_name}: {e}")

    def main(self):
        self.check_and_regenerate()